package service

import (
	"context"
	"fmt"
	chromago "github.com/amikos-tech/chroma-go"
	"github.com/amikos-tech/chroma-go/collection"
	collama "github.com/amikos-tech/chroma-go/ollama"
	"github.com/amikos-tech/chroma-go/types"
	"github.com/tmc/langchaingo/documentloaders"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/textsplitter"
	"lib/id"
	"lib/structs/enum"
	"lib/structs/lms"
	"lib/structs/tables"
	"lib/utils"
	"lms/global"
	"lms/repository"
	"os"
	"strings"
	"time"
)

var CoursewareService = new(coursewareService)

type coursewareService struct {
}

// FindDoc 查找引用的文档
func (s coursewareService) FindDoc(knowledgeId int64, que string, nResults int32) (li []lms.DocRefer, err error) {
	client, err := chromago.NewClient(global.LmsConf.RAGConfig.ChromaDBUrl)
	if err != nil {
		return
	}
	ef, err := collama.NewOllamaEmbeddingFunction(
		collama.WithBaseURL(global.LmsConf.RAGConfig.OllamaUrl),
		collama.WithModel("milkey/m3e"), // milkey/m3e，应该跟随知识库
	)
	coll, err := client.GetCollection(
		context.Background(),
		fmt.Sprintf("%d", knowledgeId),
		ef,
	)
	if err != nil {
		return
	}
	rs, err := coll.Query(context.Background(), []string{que},
		nResults, nil, nil, nil)
	if err != nil {
		return
	}
	for _, doc := range rs.Documents {
		li = append(li, lms.DocRefer{
			Content: doc[0],
		})
	}
	for i, meta := range rs.Metadatas {
		li[i].MetaData = meta[0]
	}
	return
}

// EmbeddingCourseware 嵌入课件。只接受lms课件，其他模块接入的知识库请另外写接口，不要和这个复用
func (s coursewareService) EmbeddingCourseware(ctx context.Context, li []tables.Courseware, know tables.Knowledge) (err error) {
	for _, item := range li {
		// 1. 切分文档
		docs, err := s.splitPdf(ctx, item.FilePath, know.ChunkSize, know.ChunkOverlap)
		if err != nil {
			continue
		}
		// 2. 向量化存储,
		err = s.embeddingToChroma(ctx, item, docs, know.EmbeddingModel, fmt.Sprintf("%d", know.Id))
		if err != nil {
			item.VectorStatus = enum.FileVectorStatus.Err
		} else {
			item.VectorStatus = enum.FileVectorStatus.Ok
		}
		// 3. 修改数据库中的文件状态
		if err = repository.CoursewareDb.Save(&item); err != nil {
			continue
		}
		err = repository.KnowledgeDB.BindDoc(&tables.MapKnowledgeDoc{
			KnowledgeId:  know.Id,
			DocType:      1,
			DocId:        item.Id,
			Path:         item.FilePath,
			VectorStatus: enum.FileVectorStatus.Ok,
		})
		if err != nil {
			continue
		}
	}
	return
}

// SplitPdf 对pdf文档进行分块，fp-文件路径，chunkSize-切分后每一块文件的大小，chunkOverlap-重叠数据块大小
func (s coursewareService) splitPdf(ctx context.Context, fp string, chunkSize, chunkOverlap int) (docs []schema.Document, err error) {
	f, err := os.Open(fp)
	if err != nil {
		return
	}
	finfo, err := f.Stat()
	p := documentloaders.NewPDF(f, finfo.Size())
	split := textsplitter.NewRecursiveCharacter()
	split.ChunkSize = chunkSize       // size of the chunk is number of characters
	split.ChunkOverlap = chunkOverlap // overlap is the number of characters that the chunks overlap
	docs, err = p.LoadAndSplit(ctx, split)
	return
}

// embeddingToChroma 嵌入到chromaDB，emn-embedding mode name,vName-vector collection name
// 目前看来chromaDB的功能效果好，但是没有交互ui。而且chroma基于sqlite实现存储，实际使用有待考证
func (s coursewareService) embeddingToChroma(ctx context.Context, fileInfo tables.Courseware, docs []schema.Document, emn string, vName string) (err error) {
	client, err := chromago.NewClient(global.LmsConf.RAGConfig.ChromaDBUrl)
	if err != nil {
		return err
	}
	ef, err := collama.NewOllamaEmbeddingFunction(
		collama.WithBaseURL(global.LmsConf.RAGConfig.OllamaUrl),
		collama.WithModel(emn),
	)
	li, err := client.ListCollections(ctx)
	if err != nil {
		return err
	}
	var coll *chromago.Collection
	for i := range li {
		if li[i].Name == vName {
			coll = li[i]
		}
	}
	if coll == nil {
		coll, err = client.NewCollection(
			context.Background(),
			collection.WithEmbeddingFunction(ef),
			collection.WithName(vName),
			collection.WithHNSWDistanceFunction(types.L2), // 图索引，类似数据库的索引，后面还要具体了解这个算法
		)
		if err != nil {
			return err
		}
	}
	// 创建一个记录器准备开始记录数据
	rs, err := types.NewRecordSet(
		types.WithEmbeddingFunction(ef),
		types.WithIDGenerator(types.NewULIDGenerator()),
	)
	if err != nil {
		return err
	}
	for i := range docs {
		rs.WithRecord(
			types.WithDocument(docs[i].PageContent),
			types.WithID(fmt.Sprint(id.GetSnowId())),
			types.WithMetadata("filename", fileInfo.Name),
			types.WithMetadata("sourceId", fmt.Sprintf("%d", fileInfo.Id)),
			types.WithMetadata("sourcePath", fileInfo.FilePath),
			types.WithMetadata("sourceType", "courseware"), // 标记文件类型，将来扩展用
		)
	}
	// 对数据进行编码和校验
	_, err = rs.BuildAndValidate(ctx)
	if err != nil {
		return err
	}
	_, err = coll.AddRecords(ctx, rs)
	if err != nil {
		return err
	}
	return
}

// GetCode 获取课件code，初步标记格式为 yyyyMMdd-A-0001
func (s coursewareService) GetCode() (code string, err error) {
	s1 := time.Now().Format("20060102")
	lastCode, err := repository.CoursewareDb.FindLastCode()
	if err != nil {
		return
	}
	if lastCode == "" {
		code = s1 + "-A-0001"
		return
	}
	codes := strings.Split(lastCode, "-")
	if len(codes) != 3 {
		global.LmsLog.Error("发现无效的课件code：", lastCode)
		code = s1 + "-A-0001"
		return
	}
	index := utils.StrToInt32(codes[2])
	a := []rune(codes[1])
	if index == 9999 {
		// 将字母往前移动一位
		a[0] = a[0] + 1
		index = 1
	} else {
		index++
	}
	code = fmt.Sprintf("%s-%s-%04d", s1, string(a[0]), index)
	return
}

// GetFileKind 获取文件类型，类型为0的文件表示计划外的文件需要额外处理
func (s coursewareService) GetFileKind(suffix string) (kind int) {
	if strings.HasPrefix(suffix, ".") {
		suffix = suffix[1:]
	}
	switch suffix {
	case "mp4":
		return enum.FileType.Video
	case "pdf":
		return enum.FileType.Document
	default:
		return 0
	}
}

// GetSubCoursewareFolderIds 获取下级目录列表,包含自己
func (s coursewareService) GetSubCoursewareFolderIds(parentId int64) (ids []int64, err error) {
	all, err := repository.CoursewareDb.FindAllCoursewareFolder()
	if err != nil {
		return
	}
	li := s.getSubCoursewareFolders(parentId, all)
	for _, v := range li {
		ids = append(ids, v.FolderId)
	}
	ids = append(ids, parentId)
	return
}

// getSubCoursewareFolders 递归获取下级目录,不包含自己
func (s coursewareService) getSubCoursewareFolders(parentId int64, all []tables.CoursewareFolder) (li []tables.CoursewareFolder) {
	for _, v := range all {
		if v.ParentId == parentId {
			li = append(li, v)
			sub := s.getSubCoursewareFolders(v.FolderId, all)
			li = append(li, sub...)
		}
	}
	return
}

// GetUri 获取资源访问的uri,后续的 CDN转换也使用这个方法
func (s courseService) GetUri(cw tables.Courseware) (uri string) {
	uri = "/api/lms/static/" + strings.ReplaceAll(cw.FilePath, "\\", "/")
	return
}
